import platform

from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, datasets
from setting import batch_size, cpu_workers, data_dir

#         self,
#         root: str,
#         train: bool = True,
#         transform: Optional[Callable] = None,
#         target_transform: Optional[Callable] = None,
#         download: bool = False,

train_dataset = datasets.MNIST(root=data_dir, train=True, transform=transforms.ToTensor(), download=True)
train_dataloader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True, drop_last=True,
                              num_workers=cpu_workers
                              )

test_dataset = datasets.MNIST(root=data_dir, train=False, transform=transforms.ToTensor(), download=True)
test_dataloader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False, drop_last=False,
                             num_workers=cpu_workers
                             )

if __name__ == '__main__':
    x, y = next(iter(train_dataloader))
    print(x, y.shape)
    # x, y = next(iter(test_dataloader))
    # print(x.shape, y.shape)
